# -*- coding: utf-8 -*- 
"""
Project: AppleSnApi
Creator: SF
Create time: 2025-06-26 09:36
IDE: PyCharm
Introduction:
"""
import requests
import time
import json
import concurrent.futures
from datetime import datetime
import statistics
import string
import random


def generate_random_string(length=10):
    """生成指定长度的随机字母数字字符串"""
    # 所有可能的字符：小写字母、大写字母和数字
    characters = string.ascii_letters + string.digits

    # 使用random.choices从characters中随机选择length个字符
    random_string = ''.join(random.choices(characters, k=length))

    return random_string.upper()


class APIStressTester:
    def __init__(self, url, method='GET', headers=None, data=None,
                 concurrency=10, total_requests=100, timeout=10):
        """初始化压力测试参数"""
        self.url = url
        self.method = method.upper()
        self.headers = headers or {}
        self.data = data or {}
        self.concurrency = concurrency
        self.total_requests = total_requests
        self.timeout = timeout

        # 结果统计
        self.results = []
        self.success_count = 0
        self.fail_count = 0
        self.start_time = None
        self.end_time = None

    def _make_request(self):
        """执行单个请求并返回结果"""
        start_time = time.time()
        try:
            if self.method == 'GET':
                response = requests.get(
                    self.url,
                    headers=self.headers,
                    params=self.data,
                    timeout=self.timeout
                )
            elif self.method == 'POST':
                response = requests.post(
                    self.url,
                    headers=self.headers,
                    json=self.data,
                    timeout=self.timeout
                )
            elif self.method == 'PUT':
                response = requests.put(
                    self.url,
                    headers=self.headers,
                    json=self.data,
                    timeout=self.timeout
                )
            elif self.method == 'DELETE':
                response = requests.delete(
                    self.url,
                    headers=self.headers,
                    timeout=self.timeout
                )
            else:
                raise ValueError(f"不支持的HTTP方法: {self.method}")

            # 检查响应状态码
            is_success = 200 <= response.status_code < 300
            response_time = time.time() - start_time

            return {
                'success': is_success,
                'status_code': response.status_code,
                'response_time': response_time,
                'start_time': start_time,
                'end_time': time.time(),
                'response_size': len(response.content)
            }

        except Exception as e:
            response_time = time.time() - start_time
            return {
                'success': False,
                'status_code': None,
                'response_time': response_time,
                'start_time': start_time,
                'end_time': time.time(),
                'error': str(e)
            }

    def run_test(self):
        """运行压力测试"""
        print(f"开始对 {self.url} 进行压力测试...")
        print(f"配置: 并发数={self.concurrency}, 总请求数={self.total_requests}, 超时={self.timeout}秒")

        self.start_time = time.time()

        # 使用线程池执行并发请求
        with concurrent.futures.ThreadPoolExecutor(max_workers=self.concurrency) as executor:
            # 提交所有请求任务
            future_to_request = {executor.submit(self._make_request): i for i in range(self.total_requests)}

            # 获取结果
            for i, future in enumerate(concurrent.futures.as_completed(future_to_request)):
                result = future.result()
                self.results.append(result)

                if result['success']:
                    self.success_count += 1
                else:
                    self.fail_count += 1

                # 每10%显示一次进度
                if (i + 1) % (self.total_requests // 10) == 0:
                    progress = (i + 1) / self.total_requests * 100
                    print(f"进度: {progress:.1f}% ({i + 1}/{self.total_requests})")

        self.end_time = time.time()
        self._generate_report()

    def _generate_report(self):
        """生成测试报告"""
        total_time = self.end_time - self.start_time

        # 计算响应时间统计数据
        response_times = [r['response_time'] for r in self.results]
        avg_response_time = statistics.mean(response_times)
        min_response_time = min(response_times)
        max_response_time = max(response_times)

        # 计算百分位数
        response_times.sort()
        p50 = response_times[int(len(response_times) * 0.5)]
        p90 = response_times[int(len(response_times) * 0.9)]
        p95 = response_times[int(len(response_times) * 0.95)]
        p99 = response_times[int(len(response_times) * 0.99)]

        # 计算吞吐量 (请求/秒)
        rps = self.total_requests / total_time

        # 打印报告
        print("\n" + "=" * 50)
        print(f"压力测试报告 - {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}")
        print("=" * 50)
        print(f"测试URL: {self.url}")
        print(f"HTTP方法: {self.method}")
        print(f"总请求数: {self.total_requests}")
        print(f"并发数: {self.concurrency}")
        print(f"总耗时: {total_time:.2f} 秒")
        print(f"成功请求: {self.success_count} ({self.success_count / self.total_requests * 100:.2f}%)")
        print(f"失败请求: {self.fail_count} ({self.fail_count / self.total_requests * 100:.2f}%)")
        print(f"平均响应时间: {avg_response_time * 1000:.2f} ms")
        print(f"最小响应时间: {min_response_time * 1000:.2f} ms")
        print(f"最大响应时间: {max_response_time * 1000:.2f} ms")
        print(f"响应时间中位数(p50): {p50 * 1000:.2f} ms")
        print(f"响应时间p90: {p90 * 1000:.2f} ms")
        print(f"响应时间p95: {p95 * 1000:.2f} ms")
        print(f"响应时间p99: {p99 * 1000:.2f} ms")
        print(f"吞吐量(RPS): {rps:.2f} 请求/秒")
        print("=" * 50)

        # 保存详细结果到文件
        self._save_results()

    def _save_results(self):
        """保存详细测试结果到文件"""
        timestamp = datetime.now().strftime('%Y%m%d_%H%M%S')
        filename = f"stress_test_results_{timestamp}.json"

        with open(filename, 'w') as f:
            json.dump({
                'test_config': {
                    'url': self.url,
                    'method': self.method,
                    'headers': self.headers,
                    'data': self.data,
                    'concurrency': self.concurrency,
                    'total_requests': self.total_requests,
                    'timeout': self.timeout
                },
                'test_results': self.results,
                'summary': {
                    'total_time': self.end_time - self.start_time,
                    'success_count': self.success_count,
                    'fail_count': self.fail_count,
                    'success_rate': self.success_count / self.total_requests if self.total_requests > 0 else 0
                }
            }, f, indent=4)

        print(f"详细测试结果已保存到文件: {filename}")


def main(url: str, method: str, headers: dict, body: dict, concurrency, total_requests, timeout=10):
    """
    :param url: 测试的API URL
    :param method: ['GET', 'POST', 'PUT', 'DELETE']
    :param headers: 请求头，JSON格式字符串
    :param body: 请求数据，JSON格式字符串
    :param concurrency: 并发数
    :param total_requests: 总请求数
    :param timeout: 请求超时时间(秒)
    :return:
    """
    # 创建并运行测试
    tester = APIStressTester(
        url=url,
        method=method.upper(),
        headers=headers,
        data=body if body else None,
        concurrency=concurrency,
        total_requests=total_requests,
        timeout=timeout
    )

    tester.run_test()


if __name__ == "__main__":
    test_params = {
        "url": "https://xxxx.xxx.xxx/trans/sys/proxy?xtoken=7d7da4a9fd6342e98f651d3495cc0a45&__group=sn&__method=api",
        "method": "post".upper(),
        "headers": {"Content-Type": "application/json"},
        "body": {
            "bs_id": "AAABl5uofa1P72Mt8MJDRKG0fSEM9VDz",
            "sn": "CWWLHFJDX9"
        },
        "concurrency": 30,
        "total_requests": 1000,
        "timeout": 10
    }

    main(**test_params)
